import QUANTAXIS as QA

import pandas as pd


month_data = pd.date_range(
    '1/1/1996',
    '12/31/2023',
    freq='M'
).astype(str).tolist()
start_date = '2005-03-01'
end_date = '2020-01-31'
data300 = QA.QA_fetch_index_day_adv('000300', start_date, QA.QA_util_get_recent_months(end_date, months=2))
series = pd.Series(data=month_data, index=pd.to_datetime(month_data), name='date')
timerange = series.loc[start_date: end_date].tolist()
data = {}
data['date'] = []
data['300_rtn'] = []
data['300_f_rtn'] = []
for date in timerange:
    date_next_month_last = QA.QA_util_get_real_date(QA.QA_util_get_recent_months(date, 1), towards=-1)
    date_trade = QA.QA_util_get_real_date(date, towards=-1)
    date_next_month_first = QA.QA_util_get_real_date(QA.QA_util_get_month_day_1(date_next_month_last), towards=1)
    date_next_next_month_first = QA.QA_util_get_real_date(QA.QA_util_get_month_day_1(QA.QA_util_get_recent_months(date, 2)), towards=1)

    date_trade_datetime = pd.to_datetime(date_trade)
    date_trade_next_month_datetime = pd.to_datetime(date_next_month_last)
    date_trade_next_month_first_datetime = pd.to_datetime(date_next_month_first)
    date_trade_next_next_month_first_datetime = pd.to_datetime(date_next_next_month_first)

    stock_today = data300.select_day(date_trade_datetime)
    stock_next_month_last = data300.select_day(date_trade_next_month_datetime)
    stock_next_month_first = data300.select_day(date_trade_next_month_first_datetime)
    stock_next_next_month_first = data300.select_day(date_trade_next_next_month_first_datetime)

    forward_return = (stock_next_next_month_first.CLOSE - stock_next_month_first.CLOSE) / stock_next_month_first.CLOSE
    label_return = (stock_next_month_last.CLOSE - stock_today.CLOSE) / stock_today.CLOSE
    data['date'].append(date)
    data['300_rtn'].append(label_return.values[0])
    data['300_f_rtn'].append(forward_return.values[0])  # 以1日close为基准
df = pd.DataFrame(data)
df.to_csv('hs300_return.csv', index=False)
